#!/usr/bin/env python
"""
Simple message-size benchmark
"""
import os
import sys
import time
#
from amqplib import client_0_8 as amqp

EXCHANGE = 'msgsize_ex'
QUEUE = 'msgsize_q'
OUTPUT_HDR = "run,size,dur,count"

g_exp = 0

def usage():
    prog = os.path.basename(sys.argv[0])
    sys.stderr.write("usage: %s MODE HOST COUNT\n"
                     "MODE = c for consumer, p for producer\n"
                     "HOST = rabbitmq host\n" 
                     "COUNT = number of messages\n"
                     % prog)

def writerow(sz, tm, count):
    print "%d,%d,%f,%d" % (g_exp, sz, tm, count)

def connect(host):
    conn = amqp.Connection(host=host)
    channel = conn.channel()
    channel.exchange_declare(exchange=EXCHANGE, type='direct',
                             durable=True, auto_delete=False)
    channel.queue_declare(queue=QUEUE, durable=True, auto_delete=False)
    channel.queue_bind(queue=QUEUE, exchange=EXCHANGE, routing_key=QUEUE)
    return channel

def producer(channel, n, size=1):
    buf = chr(0) * size
    for i in xrange(n):
        msg = amqp.Message(buf, delivery_mode=2)
        channel.basic_publish(msg, exchange=EXCHANGE, routing_key=QUEUE,
                              mandatory=True)

class Counter:
    def __init__(self, channel, remain):
        self.remain = remain
        self.channel = channel

    def got_msg(self, m):
        self.channel.basic_ack(m.delivery_tag)
        self.remain -= 1
            
def consumer(channel, n):
    counter = Counter(channel, n)
    tag = channel.basic_consume(queue=QUEUE, callback=counter.got_msg)
    while counter.remain > 0:
        channel.wait()
    channel.basic_cancel(tag)

def main():
    global g_exp
    try:
        mode = sys.argv[1].lower()
        host = sys.argv[2]
        n = int(sys.argv[3])
    except (IndexError, ValueError), err:
        usage()
        return 1
    powers_of_ten = range(2,7)
    sizes = (512, 1024, 2048, 4096, 8192, 16384, 32767, 65536,
             262144, 524288, 1048576, 2097152)
    for g_exp in range(1,6):
        channel = connect(host)
        if mode == 'p':
            if g_exp == 1:
                print OUTPUT_HDR
            for sz in sizes:
                t0 = time.time()
                producer(channel, n, size=sz)
                t1 = time.time()
                g = t1 - t0
                writerow(sz, g, n)
        elif mode == 'c':  
            consumer(channel, n * len(sizes))
        else:
            usage()
            return 2
        channel.close()
    return 0

if __name__ == '__main__':
   sys.exit(main())


